import cv2
import matplotlib.pyplot as plt
import numpy as np
import convertHSI


def make_histogram(img: np.ndarray, return_numpy=False):
    """
    为图像绘制直方图,若输入彩色图像则转换为HSI通道后为I通道绘制直方图
    :param img: 数据类型numpy ndarray，范围0-255 uint8，图像形状H*W*3，顺序RGB，或H*W灰度图
    :param return_numpy: 是否返回numpy格式图像
    :return: hist:list 输出一个256元素的列表表示每种灰度值出现的频率
    """
    if len(img.shape) == 3 and img.shape[2] == 3:
        (img_h, img_w, img_c) = img.shape
        _, _, i = convertHSI.rgb2hsi(img)
        img = (np.clip(i, 0, 1) * 255).astype(np.uint8)
    else:
        (img_h, img_w) = img.shape
    num_pixel = img_h * img_w
    hist = [0] * 256
    for i in range(img_h):
        for j in range(img_w):
            hist[img[i, j]] += 1 / num_pixel

    if return_numpy:
        fig = plt.figure()
        plt.bar(list(range(256)), hist)
        plt.show()
        w, h = fig.canvas.get_width_height()
        buf_ndarray = np.frombuffer(fig.canvas.tostring_rgb(), dtype='u1')
        im = buf_ndarray.reshape(h, w, 3)
        return im
    else:
        return hist


def histogram_equalization(src: np.ndarray):
    """
    灰度图像直方图均衡化
    输入张量必须为灰度，H*W*3(三个通道相同),H*W*1或H*W，0-255 uint8
    输出直方图均衡化后的灰度图 H*W, 0-255 uint8
    """
    flag = 0
    if len(src.shape) == 3 and src.shape[2] == 3:
        (img_h, img_w, img_c) = src.shape
        src_h, src_s, src_i = convertHSI.rgb2hsi(src)
        src = (src_i * 255).astype(np.uint8)
        flag = 1
    else:
        (img_h, img_w) = src.shape
    # 计算原图像直方图
    hist = make_histogram(src)
    # 累积直方图
    sum_hist = [hist[0], ]
    for i in range(1, 256):
        sum_hist.append(sum_hist[i - 1] + hist[i])
    # 计算每种像素值对应的、均衡化后的像素值
    new_hist = [0] * 256
    for i in range(256):
        new_hist[i] = round(255 * sum_hist[i])
    # 创建新图片改变像素值
    res = np.zeros_like(src)
    for i in range(img_h):
        for j in range(img_w):
            res[i, j] = new_hist[src[i, j]]
    if flag == 1:
        new_i = res / 255
        res = convertHSI.hsi2rgb(np.stack((src_h, src_s, new_i), axis=2))
    else:
        res.astype(np.uint8)
    return res


def global_linear_transform(src):  # 将灰度范围设为0~255
    """
    全局灰度线性变换
    输入张量必须为灰度，H*W*3(三个通道相同),H*W*1或H*W，0-255 uint8
    输出全局灰度拉伸后的结果，与原图像维度相同
    """
    src = src.copy().astype(np.float)
    res = (src - src.min()) / (src.max() - src.min())
    res = (res * 255).astype(np.uint8)
    return res


if __name__ == '__main__':
    image0 = cv2.imread("4.jpg", 0)
    plt.figure()
    plt.subplot(3, 2, 1)
    # 使用matplotlib中的imshow显示图像，注意参数的含义，不加参数试试
    plt.imshow(image0, vmin=0, vmax=255, cmap=plt.cm.gray)
    plt.title('original image')
    image_hist0 = make_histogram(image0)
    plt.subplot(3, 2, 2)
    plt.plot(image_hist0)

    image1 = global_linear_transform(image0)
    plt.subplot(3, 2, 3)
    plt.title('stretched image')
    plt.imshow(image1, vmin=0, vmax=255, cmap=plt.cm.gray)
    image_hist1 = make_histogram(image1)  # 统计变换后图像的各灰度值像素的个数
    plt.subplot(3, 2, 4)
    plt.plot(image_hist1)

    image2 = histogram_equalization(image0)
    plt.subplot(3, 2, 5)
    plt.imshow(image2, vmin=0, vmax=255, cmap=plt.cm.gray)
    plt.title('equalized image')
    image_hist2 = make_histogram(image2)
    plt.subplot(3, 2, 6)
    plt.plot(image_hist2)

    plt.show()
